import abc
from collections.abc import Iterator
from typing import Any, TypeVar

AnySample = TypeVar('AnySample')


class DatasetIterator(Iterator[AnySample], abc.ABC):
    """ The iterator of a Dataset

    Definitely, it supports:
        - Getting the next sample of the iteration.

    Optionally, it supports:
        - Skipping a number of samples efficiently.
        - Dumping the state of iteration for future resuming.
    """

    @abc.abstractmethod
    def __next__(self) -> AnySample:
        raise NotImplementedError

    def __skip__(self, samples_n: int):
        if samples_n == 0:
            return
        raise TypeError(f"{self!r} does not support skipping!")

    def __pause__(self) -> Any:
        raise TypeError(f"{self!r} does not support pausing!")


class SkipExhausted(StopIteration):
    def __init__(self, skipped_samples_n: int):
        self._skipped_samples_n = skipped_samples_n

    @property
    def skipped_samples_n(self) -> int:
        return self._skipped_samples_n


def skip(iterator: Iterator[AnySample], samples_n: int, *, inefficient: bool = False):
    if samples_n == 0:
        return
    if isinstance(iterator, DatasetIterator):
        return iterator.__skip__(samples_n)
    if inefficient:
        skipped = 0
        while skipped < samples_n:
            try:
                next(iterator)
            except StopIteration:
                raise SkipExhausted(skipped)
            skipped += 1
    raise TypeError(f"{iterator!r} does not support skipping!")


def pause(iterator: Iterator[AnySample]) -> Any:
    if isinstance(iterator, DatasetIterator):
        return iterator.__pause__()
    raise TypeError(f"{iterator!r} does not support pausing!")


def pause_or_error(iterator: Iterator[AnySample]) -> tuple[Any, Any]:
    try:
        return pause(iterator), None
    except (TypeError, NotImplementedError) as error:
        return None, error
